import numpy as np
from CoolProp.CoolProp import PropsSI
from enum import Enum

class PhaseType(Enum):
    LIQUID = "liquid"
    VAPOR = "vapor"

class PhysicsModel:
    """两相流物理模型类"""
    def __init__(self):
        # 封闭关系常数
        self.BETA = 0.065      # 湍流交混系数
        self.FTM = 0.8         # 端流动量因子
        self.KIJ = 0.5         # 间隙横向流阻
        
        # 重力加速度
        self.g = 9.81          # m/s^2
        
        # 固定温度（常温）
        self.T_ref = 298.15    # K (25°C)
        
        # 物性缓存
        self._properties_cache = {}
        self._pressure_range = (0.5e5, 2e5)  # 压力范围：0.5-2 bar
        
    def _get_cache_key(self, P, phase):
        """生成缓存键，按kPa取整"""
        P_kpa = round(P / 1000) * 1000
        return (P_kpa, phase)
    
    def calculate_properties(self, P, phase):
        """计算流体物性（常温常压下使用常数）
        Args:
            P: 压力 [Pa]
            phase: 相态 (液相或气相)
        Returns:
            properties: 包含密度、粘度的字典
        """
        if phase == PhaseType.LIQUID:
            # 水在常温常压下的物性
            return {
                'density': 997.0,      # kg/m³
                'viscosity': 8.9e-4    # Pa·s
            }
        else:  # VAPOR (空气)
            # 空气在常温常压下的物性
            return {
                'density': 1.225,      # kg/m³
                'viscosity': 1.81e-5   # Pa·s
            }
    
    def precalculate_properties(self, P_min, P_max, n_points=10):
        """预计算一定压力范围内的物性"""
        print("\n预计算物性:")
        pressures = np.linspace(P_min, P_max, n_points)
        for P in pressures:
            # 计算并缓存物性
            props_l = self.calculate_properties(P, PhaseType.LIQUID)
            props_g = self.calculate_properties(P, PhaseType.VAPOR)
            
            # 每隔几个点输出一次物性值
            if np.random.random() < 0.3:  # 随机输出约30%的点
                print(f"P={P/1e6:.3f} MPa:")
                print(f"  液相: ρ={props_l['density']:.1f} kg/m³, "
                      f"μ={props_l['viscosity']*1e6:.1f} μPa·s")
                print(f"  气相: ρ={props_g['density']:.1f} kg/m³, "
                      f"μ={props_g['viscosity']*1e6:.1f} μPa·s")
    
    def clear_cache(self):
        """清空物性缓存"""
        self._properties_cache.clear()
        
    def calculate_friction_factor(self, Re):
        """计算壁面摩擦因子
        Args:
            Re: Reynolds数
        Returns:
            f: 摩擦因子
        """
        # 添加保护，防止Re为0或过小
        Re = max(abs(Re), 1e-6)
        
        if Re < 2300:  # 层流
            return 64.0 / Re
        else:  # 湍流，使用Blasius公式
            return 0.184 * Re**(-0.2)
            
    def calculate_mass_flux(self, alpha, rho_g, rho_l, v_g, v_l):
        """计算质量流量
        Args:
            alpha: 空泡份额
            rho_g: 气相密度
            rho_l: 液相密度
            v_g: 气相速度
            v_l: 液相速度
        Returns:
            G_g, G_l: 气相和液相的质量流量
        """
        G_g = alpha * rho_g * v_g
        G_l = (1 - alpha) * rho_l * v_l
        return G_g, G_l
    
    def calculate_momentum_flux(self, alpha, rho_g, rho_l, v_g, v_l):
        """计算动量流量
        Args:
            alpha: 空泡份额
            rho_g: 气相密度
            rho_l: 液相密度
            v_g: 气相速度
            v_l: 液相速度
        Returns:
            M_g, M_l: 气相和液相的动量流量
        """
        M_g = alpha * rho_g * v_g * v_g
        M_l = (1 - alpha) * rho_l * v_l * v_l
        return M_g, M_l
    
    def calculate_wall_friction(self, phase_type, rho, v, D_h, mu):
        """计算壁面摩擦力
        Args:
            phase_type: 相态
            rho: 密度
            v: 速度
            D_h: 水力直径
            mu: 粘度
        Returns:
            F_w: 壁面摩擦力
        """
        # 添加保护，防止物性值为0
        rho = max(rho, 1e-6)
        mu = max(mu, 1e-10)
        D_h = max(D_h, 1e-6)
        
        Re = abs(rho * v * D_h / mu)
        f = self.calculate_friction_factor(Re)
        F_w = -0.5 * f * rho * v * abs(v) / D_h
        return F_w
    
    def calculate_interfacial_friction(self, alpha, rho_g, rho_l, v_g, v_l, D_h):
        """计算界面摩擦力
        Args:
            alpha: 空泡份额
            rho_g: 气相密度
            rho_l: 液相密度
            v_g: 气相速度
            v_l: 液相速度
            D_h: 水力直径
        Returns:
            F_i: 界面摩擦力
        """
        # 添加保护，防止空泡份额为0或1
        alpha = np.clip(alpha, 1e-6, 1-1e-6)
        rho_l = max(rho_l, 1e-6)
        D_h = max(D_h, 1e-6)
        
        v_r = v_g - v_l  # 相对速度
        C_i = 0.5  # 界面摩擦系数
        F_i = -C_i * alpha * (1 - alpha) * rho_l * v_r * abs(v_r) / D_h
        return F_i
    
    def calculate_lateral_flow_resistance(self, rho, v, s):
        """计算横向流动阻力
        Args:
            rho: 密度
            v: 横向速度
            s: 间隙宽度
        Returns:
            R: 横向流动阻力
        """
        R = -0.5 * self.KIJ * rho * v * abs(v) / s
        return R
    
    def calculate_turbulent_mixing(self, beta, s, rho_1, rho_2, v_1, v_2):
        """计算湍流交混项
        Args:
            beta: 湍流交混系数
            s: 间隙宽度
            rho_1, rho_2: 相邻子通道的密度
            v_1, v_2: 相邻子通道的速度
        Returns:
            w_t: 湍流交混流量
        """
        rho_avg = 0.5 * (rho_1 + rho_2)
        v_avg = 0.5 * abs(v_1 + v_2)
        w_t = self.BETA * s * rho_avg * v_avg
        return w_t 
    
    def calculate_drift_velocity(self, rho_g, rho_l, sigma=0.072):
        """计算空隙漂移速度
        Args:
            rho_g: 气相密度
            rho_l: 液相密度
            sigma: 表面张力系数（水在25°C下的值）
        Returns:
            v_gj: 空隙漂移速度
        """
        # Zuber-Findlay漂移通量模型
        C0 = 1.2  # 分布参数
        g = self.g  # 重力加速度
        
        # 计算漂移速度
        v_gj = 1.41 * np.sqrt(sigma * g * (rho_l - rho_g) / (rho_l * rho_l))
        return v_gj, C0

    def calculate_lateral_momentum_exchange(self, alpha, rho_g, rho_l, v_g, v_l, 
                                         v_g_adj, v_l_adj, s, gap_width):
        """计算横向动量交换
        Args:
            alpha: 空泡份额
            rho_g, rho_l: 气液相密度
            v_g, v_l: 本通道气液相速度
            v_g_adj, v_l_adj: 相邻通道气液相速度
            s: 间隙长度
            gap_width: 间隙宽度
        Returns:
            M_lat_g, M_lat_l: 气液相横向动量交换
        """
        # 计算横向速度差
        dv_g = v_g - v_g_adj
        dv_l = v_l - v_l_adj
        
        # 计算横向混合系数
        beta_g = self.BETA * alpha
        beta_l = self.BETA * (1 - alpha)
        
        # 计算横向动量交换
        M_lat_g = -beta_g * s * gap_width * rho_g * dv_g * abs(dv_g)
        M_lat_l = -beta_l * s * gap_width * rho_l * dv_l * abs(dv_l)
        
        return M_lat_g, M_lat_l

    def calculate_turbulent_mixing_momentum(self, w_t, v_1, v_2):
        """计算湍流交混动量
        Args:
            w_t: 湍流交混流量
            v_1, v_2: 相邻通道的速度
        Returns:
            M_t: 湍流交混动量
        """
        v_avg = 0.5 * (v_1 + v_2)
        M_t = w_t * v_avg
        return M_t 